# -*- coding: utf-8 -*-
r"""
Intel 场景分类 - 全量数据
HOG+颜色直方图 + 随机森林 + CV
输出文件统一加 big_ 前缀
"""
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from skimage.feature import hog
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
import joblib
import warnings
warnings.filterwarnings("ignore")


# ===================== CONFIG =====================
CODE_FILE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_ROOT = os.path.dirname(CODE_FILE_DIR)
TRAIN_DIR = os.path.join(DATA_ROOT, "seg_train")
TEST_DIR = os.path.join(DATA_ROOT, "seg_test")
CODE_DIR = os.path.join(DATA_ROOT, "codes")
RESULT_DIR = os.path.join(DATA_ROOT, "results")

IMG_SIZE = (128, 128)
HOG_WEIGHT = 0.7
COLOR_WEIGHT = 0.3
ORIENTATIONS = 9
PIXELS_PER_CELL = (8, 8)
CELLS_PER_BLOCK = (2, 2)

FAST_DEBUG = False
MAX_PER_CLASS = None

RF_PARAM_GRID = {
    "n_estimators": [400],
    "max_depth": [None],
    "min_samples_split": [2],
    "min_samples_leaf": [1],
    "n_jobs": [-1]
}
CV_FOLDS = 5

RANDOM_STATE = 42
CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
# ==================================================

os.makedirs(CODE_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

MODEL_PATH = os.path.join(RESULT_DIR, "big_best_model.pkl")   # ← 统一 big_


def extract_feature(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    hog_feat = hog(gray, orientations=ORIENTATIONS,
                   pixels_per_cell=PIXELS_PER_CELL,
                   cells_per_block=CELLS_PER_BLOCK,
                   block_norm='L2-Hys', feature_vector=True)
    color_hist = []
    for ch in range(3):
        h = cv2.calcHist([img], [ch], None, [64], [0, 256])
        h = cv2.normalize(h, h).flatten()
        color_hist.append(h)
    color_feat = np.concatenate(color_hist)
    return np.hstack([HOG_WEIGHT * hog_feat, COLOR_WEIGHT * color_feat])


def load_dataset(folder_dir):
    X, y, paths, images = [], [], [], []
    class_to_idx = {name: i for i, name in enumerate(CLASS_NAMES)}
    for cls_name in CLASS_NAMES:
        cls_dir = os.path.join(folder_dir, cls_name)
        if not os.path.isdir(cls_dir):
            continue
        fnames = [f for f in os.listdir(cls_dir)
                  if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        if FAST_DEBUG and MAX_PER_CLASS:
            fnames = fnames[:MAX_PER_CLASS]
        for fname in tqdm(fnames, desc=cls_name):
            full_path = os.path.join(cls_dir, fname)
            img = cv2.imread(full_path)
            if img is None:
                continue
            img128 = cv2.resize(img, IMG_SIZE)
            images.append(img128)
            X.append(extract_feature(img128))
            y.append(class_to_idx[cls_name])
            paths.append(full_path)
    return np.array(X, dtype=np.float32), np.array(y, dtype=np.int32), paths, images


def train():
    print("【全量】加载训练集...")
    X, y, _, _ = load_dataset(TRAIN_DIR)
    print("随机森林 GridSearchCV 调参...")
    rf = RandomForestClassifier(random_state=RANDOM_STATE)
    grid = GridSearchCV(rf, RF_PARAM_GRID, cv=CV_FOLDS,
                        scoring='accuracy', n_jobs=RF_PARAM_GRID["n_jobs"][0], verbose=2)
    grid.fit(X, y)
    print(f"最佳参数: {grid.best_params_}")
    print(f"最佳交叉验证 Accuracy: {grid.best_score_:.4f}")
    best_model = grid.best_estimator_
    joblib.dump(best_model, MODEL_PATH)
    print("模型已保存 ->", MODEL_PATH)
    return best_model


def evaluate(model):
    print("【全量】加载测试集...")
    X_test, y_test, paths, test_images = load_dataset(TEST_DIR)
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    print(f"测试集 Accuracy: {acc:.4f}")

    # 1. 测试图片数组
    np.save(os.path.join(RESULT_DIR, "big_all_test_images.npy"),
            np.array(test_images, dtype=np.uint8))
    # 2. 模型备份
    joblib.dump(model, os.path.join(RESULT_DIR, "big_best_model_full.pkl"))
    # 3. 混淆矩阵
    cm = confusion_matrix(y_test, y_pred, labels=range(len(CLASS_NAMES)))
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=CLASS_NAMES)
    disp.plot(cmap='Blues', xticks_rotation=45)
    plt.tight_layout()
    plt.savefig(os.path.join(RESULT_DIR, "big_confusion_matrix.png"), dpi=300)
    plt.close()
    # 4. 抽样图
    show_predictions(paths, y_test, y_pred, 12)
    return acc


def show_predictions(paths, y_test, y_pred, n_show=12):
    import random
    idxs = random.sample(range(len(paths)), n_show)
    plt.figure(figsize=(12, 9))
    for i, idx in enumerate(idxs, 1):
        img = cv2.imread(paths[idx])[:, :, ::-1]
        true_lab = CLASS_NAMES[y_test[idx]]
        pred_lab = CLASS_NAMES[y_pred[idx]]
        plt.subplot(3, 4, i)
        plt.imshow(img)
        plt.title(f"True: {true_lab}\nPred: {pred_lab}",
                  color=("green" if true_lab == pred_lab else "red"))
        plt.axis('off')
    plt.tight_layout()
    plt.savefig(os.path.join(RESULT_DIR, "big_pred_samples.png"), dpi=300)
    plt.close()


# ===================== 纯控制台预测 =====================
def predict_one_image(model, img_path):
    """
    输入：模型 + 图片绝对路径
    输出：预测类别字符串
    """
    img_bgr = cv2.imread(img_path)
    if img_bgr is None:
        raise ValueError("无法读取图片，请检查路径")
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    img128 = cv2.resize(img_bgr, IMG_SIZE)
    feat = extract_feature(img128).reshape(1, -1)
    pred_id = int(model.predict(feat)[0])
    return CLASS_NAMES[pred_id], img_rgb


def console_demo(model):
    """
    循环控制台交互，直到用户输入空行退出
    """
    print("\n===== 控制台预测模式 =====")
    print("（直接回车可退出）")
    while True:
        path = input("\n请输入图片绝对路径：").strip()
        if not path:
            print("已退出")
            break
        try:
            pred, img = predict_one_image(model, path)
            print(f"预测类别：{pred}")
            # matplotlib 展示原图
            plt.imshow(img)
            plt.axis('off')
            plt.title(f"Predict: {pred}")
            plt.show()
        except Exception as e:
            print("出错：", e)


# ===================== 唯一入口 =====================
if __name__ == "__main__":
    if os.path.isfile(MODEL_PATH):
        print("检测到已训练模型，加载中…")
        best_model = joblib.load(MODEL_PATH)
        console_demo(best_model)          # ← 改为控制台交互
    else:
        print("未检测到模型，进入全量训练流程…")
        best_model = train()
        evaluate(best_model)
        print("【全量】训练完成！接下来可重新运行脚本进行控制台预测。")